#include "ComputeEdgeStrength.h"

#include "../DWMRICore/Fiber.h"
#include "../DWMRICore/Scalar.h"

#include <omp.h>

enum ES_METHOD{
	ES_NONE,
	ES_FTLE,		/* FTLE */
	ES_PP,			/* Point Pair */
	ES_NEW			/* New method, compute graident based on estimated function */
};

char *fiber_pathname = NULL;
char *dst_pathname = NULL;
ES_METHOD method = ES_NONE;

char **confPathname = NULL;
char **vertPathname = NULL;
int w, h, d, scale;
float stepSize, maxLength;

float fiber_length = 0.0f;

void PrintUsage()
{
	printf("Usage: ComputeEdgeStrength -fiber <configure_file> -dst <dst_file> -method <method> -length <fiber_length>\n");
}

int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-fiber") == 0) {
			fiber_pathname = argv[++i];
		} else if (strcmp(argv[i], "-dst") == 0) {
			dst_pathname = argv[++i];
		} else if (strcmp(argv[i], "-length") == 0) {
			fiber_length = (float)atof(argv[++i]);
		} else if (strcmp(argv[i], "-method") == 0) {
            i++;
			if (strcmp(argv[i], "FTLE") == 0 || strcmp(argv[i], "ftle") == 0) {
				method = ES_FTLE;
			} else if (strcmp(argv[i], "PP") == 0 || strcmp(argv[i], "pp") == 0) {
				method = ES_PP;
			} else if (strcmp(argv[i], "NEW") == 0 || strcmp(argv[i], "new") == 0) {
				method = ES_NEW;
			} else {
				return -1;
			}
		} else {
			return -1;
		}
	}

	if (fiber_pathname == NULL || dst_pathname == NULL || method == ES_NONE)
		return -2;

	return 0;
}

void ComputeEdgeStrengthByFTLE(float *edgeStrength)
{
	CFiber *fibers[3];
	int cw, ch;
	
    int offset[9];
    offset[ 0] = 0;
    offset[ 1] = -1;
    offset[ 2] = 1;
    offset[ 3] = -w-1;
    offset[ 4] = -w;
    offset[ 5] = -w+1;
    offset[ 6] = w-1;
    offset[ 7] = w;
    offset[ 8] = w+1;
	
    omp_set_num_threads(64);

	float length = fiber_length > maxLength ? maxLength : fiber_length;
	int id = length / stepSize;
	float delta = 1.0f / (float)scale;

	for (int z = 1; z < d - 1; ++z) {
		ReadFibers(confPathname[z-1], vertPathname[z-1], &(fibers[0]), cw, ch);
		ReadFibers(confPathname[z+0], vertPathname[z+0], &(fibers[1]), cw, ch);
		ReadFibers(confPathname[z+1], vertPathname[z+1], &(fibers[2]), cw, ch);
		
		
		for (int y = 1; y < h - 1; ++y) {
#pragma omp parallel for            
            for (int x = 1; x < w - 1; ++x) {
				int index = y * w + x;
				CFiber *f[27];

				f[0] = &(fibers[1][index]);			// fiber of current voxel

				for (int i = 1; i < 9; ++i) {
					f[i] = &(fibers[1][index+offset[i]]);
				}
				for (int i = 0; i < 9; ++i) {
					f[i+9] = &(fibers[0][index+offset[i]]);
				}
				for (int i = 0; i < 9; ++i) {
					f[i+18] = &(fibers[2][index+offset[i]]);
				}

				//f[1] = &(fibers[1][index-1]);		// fibers of neighborhood voxels
				//f[2] = &(fibers[1][index+1]);
				//f[3] = &(fibers[1][index-w]);
				//f[4] = &(fibers[1][index+w]);
				//f[5] = &(fibers[0][index]);
				//f[6] = &(fibers[2][index]);
			
				edgeStrength[z*w*h+index] = EdgeStrengthMethodFTLE(f, id, delta);
			}
		}

		delete[] fibers[0];
		delete[] fibers[1];
		delete[] fibers[2];

		printf("z = %d\n", z);
	}
	
	
}

void ComputeEdgeStrengthByPointPair(float *edgeStrength)
{
	
}

/*
void ComputeEdgeStrengthByNewMethod_Test(float *edgeStrength)
{
	for (int z = 0; z < d; ++z) {
        for (int iStep = 0; iStep < 32; ++iStep) {    
            printf("%d ... ", iStep);

            CFiber **fibers = NULL;
            int *count = NULL;
            int cw, ch;

            char _confPathname[256];
            char _vertPathname[256];
            sprintf(_confPathname, "%s_%d.conf", "/home/hypnos/Project/Dataset/Alex/new_method_test/fibers/test_slice", iStep);
            sprintf(_vertPathname, "%s_%d.vert", "/home/hypnos/Project/Dataset/Alex/new_method_test/fibers/test_slice", iStep);
            
            ReadFibers(_confPathname, _vertPathname, &fibers, &count, cw, ch);
            
            for (int y = 0; y < ch; ++y) {
                for (int x = 0; x < cw; ++x) {
                    int index = y * cw + x;
                    int id = (y + iStep * 16) * cw + x;
                    
                    CFiber **f = new CFiber*[count[index]];
                    for (int k = 0; k < count[index]; ++k) {
                        f[k] = &(fibers[index][k]);
                    }
                   
                    edgeStrength[id] = EdgeStrengthMethodNEW(f, count[index]);

                    delete[] f;
                }
            }

            for (int i = 0; i < cw*ch; ++i) {
                SafeDeleteArray(fibers[i]);
            }
            SafeDeleteArray(fibers);
            SafeDeleteArray(count);

            printf("done.\n");
        }
    }
}*/


void ComputeEdgeStrengthByNewMethod(float *edgeStrength)
{
    int offset[9];
    offset[ 0] = 0;
    offset[ 1] = -1;
    offset[ 2] = 1;
    offset[ 3] = -w-1;
    offset[ 4] = -w;
    offset[ 5] = -w+1;
    offset[ 6] = w-1;
    offset[ 7] = w;
    offset[ 8] = w+1;
    
    omp_set_num_threads(16);

	for (int z = 1; z < d - 1; ++z) {
        CFiber **fibers[3];
        int *count[3];
        int cw, ch;
        
        ReadFibers(confPathname[z-1], vertPathname[z-1], &(fibers[0]), &(count[0]), cw, ch);
        ReadFibers(confPathname[z+0], vertPathname[z+0], &(fibers[1]), &(count[1]), cw, ch);
		ReadFibers(confPathname[z+1], vertPathname[z+1], &(fibers[2]), &(count[2]), cw, ch);
		
		for (int y = 1; y < h-1; ++y) {
#pragma omp parallel for            
			for (int x = 1; x < w-1; ++x) {
				int index = y * w + x;
                edgeStrength[z*w*h+index] = 0.0f;      

				//bool flags[3] = {false, false, false};
				//if (count[1][index] == 1) {
				//	flags[0] = true;
				//} else {
				//	for (int k = 0; k < count[1][index]; ++k) {
				//		CVector3F dir = fibers[1][index][k].m_seed.dir;
				//		flags[k] = true;
				//		for (int l = 0; l < 3; ++l) {
				//			for (int i = 0; i < 9; ++i) {
				//				if (l == 1 && i == 0)
				//					continue;
				//				float max_v = -1.0f;
				//				int max_index;
				//				int temp_index = index + offset[i];
				//				int size = count[l][temp_index];
				//				for (int j = 0; j < size; ++j) {
				//					float v = InnerProduct(dir, fibers[l][temp_index][j].m_seed.dir);
				//					if (fabs(v) > max_v) {
				//						max_v = fabs(v);
				//						max_index = j;
				//					}                            
				//				}
				//				if (max_v < 0.866f)
				//					flags[k] = false;
				//			}
				//		}
				//	}
				//	if (flags[0] == false && flags[1] == false && flags[2] == false) {
				//		flags[0] = true;
				//	}
				//}

				//CFiber ghostFibers[27];
				//for (int l = 0; l < 3; ++l) {
				//	for (int i = 0; i < 9; ++i) {
				//		int temp_index = index + offset[i];
				//		CPoint3F pos = fibers[l][temp_index][0].m_seed.pos;
				//		if (l == 0) {
				//			ghostFibers[9+i].CreateFiber(&pos, &pos, 1, 1, &(fibers[l][temp_index][0].m_seed));
				//		} else if (l == 1) {
				//			ghostFibers[i].CreateFiber(&pos, &pos, 1, 1, &(fibers[l][temp_index][0].m_seed));
				//		} else {
				//			ghostFibers[18+i].CreateFiber(&pos, &pos, 1, 1, &(fibers[l][temp_index][0].m_seed));
				//		}
				//	}
				//}

				bool skip_flag[3] = {false, false, false};
				//for (int k = count[1][index]; k < 3; ++k) {
				//	skip_flag[k] = true;
				//}
				//for (int k = 0; k < count[1][index]; ++k) {
				//	CVector3F dir = fibers[1][index][k].m_seed.dir;
				//	if (count[1][index] > 1 && fabs(dir.m_y) > 0.7) {
				//		skip_flag[k] = true;
				//	}

				//	//for (int l = 0; l < 3; ++l) {
				//	//	for (int i = 0; i < 9; ++i) {
				//	//		if (l == 1 && i == 0)
				//	//			continue;
				//	//			float max_v = -1.0f;
				//	//			int max_index;
				//	//			int temp_index = index + offset[i];
				//	//			int size = count[l][temp_index];
				//	//			for (int j = 0; j < size; ++j) {
				//	//				float v = InnerProduct(dir, fibers[l][temp_index][j].m_seed.dir);
				//	//				if (fabs(v) > max_v) {
				//	//					max_v = fabs(v);
				//	//					max_index = j;
				//	//				}                            
				//	//			}
				//	//			if (max_v < 0.866f)
				//	//				skip_flag[k] = true;
				//	//	}
				//	//}
				//}
				//if (skip_flag[0] == true && skip_flag[1] == true && skip_flag[2] == true)
				//	skip_flag[0] = false;
				

                for (int k = 0; k < count[1][index]; ++k) {
					if (skip_flag[k] == true)
						continue;

                    CFiber *f[27];
                    f[0] = &(fibers[1][index][k]);
                    CVector3F dir = f[0]->m_seed.dir;

                    int id = 1;
					bool flag = true;
					for (int l = 0; l < 3; ++l) {
						for (int i = 0; i < 9; ++i) {
							if (l == 1 && i == 0)
								continue;
							float max_v = -1.0f;
							int max_index;
							int temp_index = index + offset[i];
							//int size = count[1][temp_index] > 2 ? 2 : count[1][temp_index];
							int size = count[l][temp_index];
							for (int j = 0; j < size; ++j) {
								float v = InnerProduct(dir, fibers[l][temp_index][j].m_seed.dir);
								if (fabs(v) > max_v) {
									max_v = fabs(v);
									max_index = j;
									if (l == 0) {
										f[9+i] = &(fibers[l][temp_index][j]);
									} else if (l == 1) {
										f[i] = &(fibers[l][temp_index][j]);
									} else {
										f[18+i] = &(fibers[l][temp_index][j]);
									}
								}                            
							}
							//if (max_v < 0.866f) {
							//	if (l == 0) {
							//		f[9+i] = &(ghostFibers[9+i]);
							//	} else if (l == 1) {
							//		f[i] = &(ghostFibers[i]);
							//	} else {
							//		f[18+i] = &(ghostFibers[18+i]);
							//	}
							//}
						}
					}

					float tempStrength = EdgeStrengthMethodNEW(f, 27);
					if (tempStrength > edgeStrength[z*w*h+index]) {
						edgeStrength[z*w*h+index] = tempStrength;
					}
                    //edgeStrength[z*w*h+index] += EdgeStrengthMethodNEW(f, 27);
                    /*if (id > 7) {
                        edgeStrength[z*w*h+index] += EdgeStrengthMethodNEW(f, id);
                    }*/
                    //printf("x = %d, y = %d z = %d, edge strength = %f\n", x, y, z, edgeStrength[z*w*h+index]);
                    //exit(1);
                }
			}
		}

        for (int i = 0; i < w * h; ++i) {
            SafeDeleteArray(fibers[0][i]);
            SafeDeleteArray(fibers[1][i]);
            SafeDeleteArray(fibers[2][i]);
        }
	    SafeDeleteArray(fibers[0]);
        SafeDeleteArray(fibers[1]);
        SafeDeleteArray(fibers[2]);	
	    SafeDeleteArray(count[0]);
        SafeDeleteArray(count[1]);
        SafeDeleteArray(count[2]);	

        printf("z = %d\n", z);
        fflush(stdout);
	}
	
	//delete[] fibers[0];
	//delete[] fibers[1];

}

void test()
{
    CFiber fibers[7];
    CPoint3F F[2], B[2];
    Seed seed;

    // fiber 0
    F[0] = CPoint3F(4, 4, 4); F[1] = CPoint3F(5, 4, 4);
    B[0] = CPoint3F(4, 4, 4); B[1] = CPoint3F(3, 4, 4);
    seed.pos = CPoint3F(4, 4, 4); seed.dir = CVector3F(1, 0, 0);
    fibers[0].CreateFiber(F, B, 2, 2, &seed);

    // fiber 1
    F[0] = CPoint3F(4.5, 4, 4); F[1] = CPoint3F(5.5, 4, 4);
    B[0] = CPoint3F(4.5, 4, 4); B[1] = CPoint3F(3.5, 4, 4);
    seed.pos = CPoint3F(4.5, 4, 4); seed.dir = CVector3F(1, 0, 0);
    fibers[1].CreateFiber(F, B, 2, 2, &seed);

    // fiber 2
    F[0] = CPoint3F(3.5, 4, 4); F[1] = CPoint3F(2.5, 4, 4);
    B[0] = CPoint3F(3.5, 4, 4); B[1] = CPoint3F(4.5, 4, 4);
    seed.pos = CPoint3F(3.5, 4, 4); seed.dir = CVector3F(-1, 0, 0);
    fibers[2].CreateFiber(F, B, 2, 2, &seed);

    // fiber 3
    F[0] = CPoint3F(4, 4.5, 4); F[1] = CPoint3F(5, 4.5, 4);
    B[0] = CPoint3F(4, 4.5, 4); B[1] = CPoint3F(3, 4.5, 4);
    seed.pos = CPoint3F(4, 4.5, 4); seed.dir = CVector3F(1, 0, 0);
    fibers[3].CreateFiber(F, B, 2, 2, &seed);

    // fiber 4
    F[0] = CPoint3F(4, 3.5, 4); F[1] = CPoint3F(3, 3.5, 4);
    B[0] = CPoint3F(4, 3.5, 4); B[1] = CPoint3F(5, 3.5, 4);
    seed.pos = CPoint3F(4, 3.5, 4); seed.dir = CVector3F(-1, 0, 0);
    fibers[4].CreateFiber(F, B, 2, 2, &seed);

    // fiber 5
    F[0] = CPoint3F(4, 4, 4.5); F[1] = CPoint3F(3, 4, 4.5);
    B[0] = CPoint3F(4, 4, 4.5); B[1] = CPoint3F(5, 4, 4.5);
    seed.pos = CPoint3F(4, 4, 4.5); seed.dir = CVector3F(-1, 0, 0);
    fibers[5].CreateFiber(F, B, 2, 2, &seed);
 
    // fiber 6
    F[0] = CPoint3F(4, 4, 3.5); F[1] = CPoint3F(5, 4, 3.5);
    B[0] = CPoint3F(4, 4, 3.5); B[1] = CPoint3F(3, 4, 3.5);
    seed.pos = CPoint3F(4, 4, 3.5); seed.dir = CVector3F(1, 0, 0);
    fibers[6].CreateFiber(F, B, 2, 2, &seed);
    
    CFiber *f[7];
    f[0] = &(fibers[0]);
    //f[1] = &(fibers[1]);
    //f[2] = &(fibers[2]);
    f[3] = &(fibers[3]);
    f[4] = &(fibers[4]);
    f[1] = &(fibers[5]);
    f[2] = &(fibers[6]);

    float strength = EdgeStrengthMethodNEW(f, 5);
}

int main(int argc, char *argv[])
{
    //test();

	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}
	
	/* read the configure file */
	ReadConfigureFile(fiber_pathname, &confPathname, &vertPathname, w, h, d, scale, stepSize, maxLength);
	
	/* compute edge strength for each voxel */
	float *edgeStrength = new float[w*h*d];
	memset(edgeStrength, 0, sizeof(float)*w*h*d);
	
	if (method == ES_FTLE) {
		ComputeEdgeStrengthByFTLE(edgeStrength);
	} else if (method == ES_PP) {
		ComputeEdgeStrengthByPointPair(edgeStrength);
	} else if (method == ES_NEW) {
		ComputeEdgeStrengthByNewMethod(edgeStrength);
	} else {
		;
	}

    /* save the result */
    CScalar *result = new CScalar();
    result->CreateScalar(w, h, d, edgeStrength);
    result->SaveScalarFile(dst_pathname);

    delete result;
	
	for (int z = 0; z < d; ++z) {
		delete[] confPathname[z];
		delete[] vertPathname[z];
	}
	delete[] confPathname;
	delete[] vertPathname;
	
	delete[] edgeStrength;
}
